# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

import logging
from typing import Optional

import coin.flow.topic_map as topic_map
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase


class BitmexSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               subscription: Optional[str] = None,
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    subscription = (subscription or ','.join(
        ['orderBook10', 'orderBookL2', 'trade', 'instrument', 'liquidation']))
    self._register_url('wss://www.bitmex.com/realtime?subscribe=%s' % subscription)

    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_order_book_l2 = topic_map.search_topic_id('bitmex_ws_orderBookL2_%s%s' %
                                                    (group_id, worker_id))[0]
    topic_order_book_10 = topic_map.search_topic_id('bitmex_ws_orderBook10_%s%s' %
                                                    (group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('bitmex_ws_trade_%s%s' % (group_id, worker_id))[0]
    topic_instrument = topic_map.search_topic_id('bitmex_ws_instrument_%s' % worker_id)[0]
    topic_liquidation = topic_map.search_topic_id('bitmex_ws_liquidation_%s' % worker_id)[0]
    self._publisher = {
        'orderBookL2': queue.get_writer(topic_id=topic_order_book_l2[1]),
        'orderBook10': queue.get_writer(topic_id=topic_order_book_10[1]),
        'trade': queue.get_writer(topic_id=topic_trade[1]),
        'instrument': queue.get_writer(topic_id=topic_instrument[1]),
        'liquidation': queue.get_writer(topic_id=topic_liquidation[1])
    }

    self._enable_topic_event_publisher(
        queue,
        'bitmex_ws_topic_event',
        [topic_order_book_l2[0], topic_order_book_10[0], topic_trade[0], topic_instrument[0]])

    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def on_close(self):
    try:
      super().on_close()
    finally:
      for _, publisher in (self._publisher or {}).items():
        publisher.close()
      self._publisher = None

  def on_message(self, timestamp: int, message: dict):
    if self.stopped:
      return

    if 'table' in message:
      table = message['table']
      try:
        self._publisher[table].write(timestamp, message)
      except KeyError:
        self._logger.warning('Unknown table: %s' % table)
